from datetime import datetime, timedelta
from typing import Dict, List, Optional, Tuple
import pandas as pd
from loguru import logger
from services.data_view_plugin_service import DataViewPluginService


class IndexDataModel:
    
    def __init__(self):
        self.data_service = DataViewPluginService()

    def get_table_data(self, days_limit: int = 365) -> Optional[List]:
        return self.data_service.load_all_index_data(days_limit=days_limit)

    def get_index_history_data(self, codes: List[str], start_date: str, end_date: str) -> Dict[str, pd.DataFrame]:
        all_data_df = self.data_service.stock_index_service.get_history_data(
            codes=codes,
            start_date=start_date,
            end_date=end_date
        )

        index_data = {}
        if all_data_df is not None and not all_data_df.empty and 'code' in all_data_df.columns:
            for code, group in all_data_df.groupby('code'):
                index_data[str(code)] = group.reset_index(drop=True)
        
        for code in codes:
            if code not in index_data:
                logger.warning(f"未获取到指数 {code} 的数据")
                index_data[code] = None

        return index_data

    def calculate_ratio(self, df1: pd.DataFrame, df2: pd.DataFrame) -> pd.DataFrame:
        try:
            df1 = df1.sort_values('date').reset_index(drop=True)
            df2 = df2.sort_values('date').reset_index(drop=True)
            
            merged_df = pd.merge(df1[['date', 'close']], df2[['date', 'close']],
                               on='date', how='inner', suffixes=('_1', '_2'))
            
            merged_df['cumulative_return_1'] = merged_df['close_1'] / merged_df['close_1'].iloc[0]
            merged_df['cumulative_return_2'] = merged_df['close_2'] / merged_df['close_2'].iloc[0]
            
            merged_df['size_ratio'] = merged_df['cumulative_return_1'] / merged_df['cumulative_return_2']
            
            merged_df = merged_df[merged_df['size_ratio'] > 0]
            
            return merged_df[['date', 'size_ratio']].sort_values('date')
            
        except Exception as e:
            logger.error(f"计算比价失败: {e}")
            return pd.DataFrame()

    def calculate_cumulative_return(self, df: pd.DataFrame) -> pd.DataFrame:
        if df is None or df.empty:
            return pd.DataFrame()
        
        df = df.sort_values('date').reset_index(drop=True)
        df['cumulative_return'] = (df['close'] / df['close'].iloc[0]) * 100
        return df

    def calculate_volume_ratio(self, indices_data: Dict[str, pd.DataFrame]) -> Tuple[List[str], List[Dict]]:
        valid_dates = None
        volume_data = {}
        
        for code, df in indices_data.items():
            if df is not None and not df.empty and 'money' in df.columns:
                df = df.sort_values('date').reset_index(drop=True)
                df = df[df['money'] > 0].dropna(subset=['money'])
                if not df.empty:
                    volume_data[code] = df
                    if valid_dates is None:
                        valid_dates = set(df['date'].dt.strftime('%Y-%m-%d'))
                    else:
                        valid_dates &= set(df['date'].dt.strftime('%Y-%m-%d'))
        
        if not volume_data or not valid_dates:
            return [], []
        
        sorted_dates = sorted(list(valid_dates))
        
        daily_ratios = {}
        for code in volume_data:
            daily_ratios[code] = []
            df = volume_data[code]
            for date_str in sorted_dates:
                date_data = df[df['date'].dt.strftime('%Y-%m-%d') == date_str]
                if not date_data.empty:
                    daily_ratios[code].append(float(date_data.iloc[0]['money']))
                else:
                    daily_ratios[code].append(0.0)
        
        return sorted_dates, daily_ratios

    def get_default_date_range(self) -> Tuple[str, str]:
        end_date = datetime.now()
        start_date = end_date - timedelta(days=3*30)
        return start_date.strftime('%Y-%m-%d'), end_date.strftime('%Y-%m-%d')

    @staticmethod
    def get_index_pairs() -> Dict[str, Dict]:
        return {
            "sz50_vs_zz2000": {
                "codes": ("000016", "932000"),
                "name": "上证50/中证2000"
            },
            "csi300_vs_csi500": {
                "codes": ("000300", "000905"),
                "name": "沪深300/中证500"
            },
            "csi300_vs_csi1000": {
                "codes": ("000300", "000852"),
                "name": "沪深300/中证1000"
            },
            "csi300_vs_zz2000": {
                "codes": ("000300", "932000"),
                "name": "沪深300/中证2000"
            },
        }

    @staticmethod
    def get_cumulative_indices() -> Dict[str, str]:
        return {
            '000985': '中证全指',
            '000016': '上证50',
            '000300': '沪深300',
            '000905': '中证500',
            '000852': '中证1000',
            '932000': '中证2000',
        }

    @staticmethod
    def get_volume_ratio_indices() -> Dict[str, str]:
        return {
            '000300': '沪深300',
            '000905': '中证500', 
            '000852': '中证1000',
            '932000': '中证2000',
        }

    @staticmethod
    def get_style_indices() -> Dict[str, str]:
        return {
            "399373": '大盘价值',
            '399372': '大盘成长',
            '399375': '中盘价值',
            '399374': '中盘成长',
            '399377': '小盘价值',
            '399376': '小盘成长'
        }

    @staticmethod
    def get_asset_class_indices() -> Dict[str, str]:
        return {
            '000985': '中证全指',
            '000922': '中证红利',
            '000012': "国债指数",
            '000832': '中证转债',
            '000066': '上证商品',
        }

    @staticmethod
    def get_sector_indices() -> Dict[str, str]:
        return {
            '000932': '中证消费',
            '000688': '科创50',
            '000933': '中证医药',
            '000934': '中证金融',
            '399997': '中证白酒',
            '000006': '地产指数',
            '000928': '中证能源',
            '931151': '光伏产业',
            '980017': '国证芯片',
            '399986': '中证银行',
        }
